package main

import (
	"database/sql"
	"fmt"
	"github.com/coredns/coredns/plugin"
	"github.com/coredns/coredns/request"
	_ "github.com/mattn/go-sqlite3"
	"github.com/miekg/dns"
	"golang.org/x/net/context"
	"os"
	"strings"
	"time"
)

type ItDns struct {
	Next               plugin.Handler
	Dsn                string
	TablePrefix        string
	MaxLifetime        time.Duration
	MaxOpenConnections int
	MaxIdleConnections int
	Ttl                uint32
	tableName          string
	lastZoneUpdate     time.Time
	zoneUpdateTime     time.Duration
	zones              []string
}

const Name = "itdns"

// ServeDNS implements the plugin.Handler interface.
func (itDns *ItDns) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg) (int, error) {

	log.Info("执行handler --> ServeDNS")

	state := request.Request{W: w, Req: r}

	qName := state.Name()
	qType := state.Type()

	if strings.HasSuffix(qName, ".") {
		qName = qName[:len(qName)-1]
	}

	if time.Since(itDns.lastZoneUpdate) > itDns.zoneUpdateTime {
		err := itDns.loadZones()
		if err != nil {
			return itDns.errorResponse(state, dns.RcodeServerFailure, err)
		}
	}

	qZone := plugin.Zones(itDns.zones).Matches(qName)
	if qZone == "" {
		return plugin.NextOrFailure(itDns.Name(), itDns.Next, ctx, w, r)
	}

	records, err := itDns.findRecord(qZone, qName, qType)

	log.Info("查询结果:", records, err)

	if err != nil {
		return itDns.errorResponse(state, dns.RcodeServerFailure, err)
	}

	if len(records) == 0 {
		// no record found but we are going to return a SOA
		rec := &Record{
			Name:  qName,
			Type:  "SOA",
			Value: "{}",
		}
		records = append(records, rec)
	}

	if qType == "AXFR" {
		return itDns.errorResponse(state, dns.RcodeNotImplemented, nil)
	}

	answers := make([]dns.RR, 0, 10)
	extras := make([]dns.RR, 0, 10)

	for _, record := range records {
		var answer dns.RR
		switch record.Type {
		case "A":
			answer, extras, err = record.AsARecord()
		case "AAAA":
			answer, extras, err = record.AsAAAARecord()
		case "CNAME":
			answer, extras, err = record.AsCNAMERecord()
		case "SOA":
			answer, extras, err = record.AsSOARecord()
		case "SRV":
			answer, extras, err = record.AsSRVRecord()
		case "NS":
			answer, extras, err = record.AsNSRecord()
		case "MX":
			answer, extras, err = record.AsMXRecord()
		case "TXT":
			answer, extras, err = record.AsTXTRecord()
		case "CAA":
			answer, extras, err = record.AsCAARecord()
		default:
			return itDns.errorResponse(state, dns.RcodeNotImplemented, nil)
		}

		if err != nil {
			return itDns.errorResponse(state, dns.RcodeServerFailure, err)
		}
		if answer != nil {
			answers = append(answers, answer)
		}
	}

	m := new(dns.Msg)
	m.SetReply(r)
	m.Authoritative = true
	m.RecursionAvailable = false
	m.Compress = true

	m.Answer = append(m.Answer, answers...)
	m.Extra = append(m.Extra, extras...)

	state.SizeAndDo(m)
	m = state.Scrub(m)
	_ = w.WriteMsg(m)
	return dns.RcodeSuccess, nil
}

// Name implements the Handler interface.
func (itDns *ItDns) Name() string { return Name }

func (itDns *ItDns) errorResponse(state request.Request, rCode int, err error) (int, error) {
	m := new(dns.Msg)
	m.SetRcode(state.Req, rCode)
	m.Authoritative, m.RecursionAvailable, m.Compress = true, false, true

	state.SizeAndDo(m)
	_ = state.W.WriteMsg(m)
	// Return success as the rCode to signal we have written to the client.
	return dns.RcodeSuccess, err
}

func (itDns *ItDns) findRecord(zone string, name string, types ...string) ([]*Record, error) {

	log.Info("执行 itdns --> findRecord")

	db, err := itDns.open()

	if err != nil {
		return nil, err
	}

	defer itDns.close(db)

	query := strings.TrimSuffix(name, "."+zone)
	sqlQuery := fmt.Sprintf("SELECT name, zone, ttl, type, value FROM %s WHERE zone = ? AND name = ? AND type IN ('%s')",
		itDns.tableName,
		strings.Join(types, "','"))

	log.Info("查询记录:", sqlQuery)

	result, err := db.Query(sqlQuery, zone, query)

	if err != nil {
		log.Error("查询错误:", err)
		return nil, err
	}

	//var recordName string
	//var recordZone string
	//var recordType string
	//var ttl uint32
	//var content string

	r := &Record{
		//Name:       recordName,
		//Zone:       recordZone,
		//Type: recordType,
		//Ttl:        ttl,
		//Value:    content,
		handler: itDns,
	}

	records := make([]*Record, 0)
	for result.Next() {
		err = result.Scan(&r.Name, &r.Zone, &r.Ttl, &r.Type, &r.Value)
		if err != nil {
			return nil, err
		}

		//if !strings.HasSuffix(r.Zone, ".") {
		//    r.Zone = r.Zone + "."
		//}

		records = append(records, r)
	}

	return records, nil
}

func (itDns *ItDns) loadZones() error {
	db, err := itDns.open()
	if err != nil {
		return err
	}
	defer itDns.close(db)

	result, err := db.Query("SELECT DISTINCT zone FROM " + itDns.tableName)
	if err != nil {
		return err
	}

	var zone string
	zones := make([]string, 0)
	for result.Next() {
		err = result.Scan(&zone)
		if err != nil {
			return err
		}

		zones = append(zones, zone)
	}

	itDns.lastZoneUpdate = time.Now()
	itDns.zones = zones

	return nil
}

func (itDns *ItDns) hosts(zone string, name string) ([]dns.RR, error) {
	recs, err := itDns.findRecord(zone, name, "A", "AAAA", "CNAME")
	if err != nil {
		return nil, err
	}

	answers := make([]dns.RR, 0)

	for _, rec := range recs {
		switch rec.Type {
		case "A":
			aRec, _, err := rec.AsARecord()
			if err != nil {
				return nil, err
			}
			answers = append(answers, aRec)
		case "AAAA":
			aRec, _, err := rec.AsAAAARecord()
			if err != nil {
				return nil, err
			}
			answers = append(answers, aRec)
		case "CNAME":
			aRec, _, err := rec.AsCNAMERecord()
			if err != nil {
				return nil, err
			}
			answers = append(answers, aRec)
		}
	}

	return answers, nil
}

func (itDns *ItDns) open() (*sql.DB, error) {

	log.Info("初始化数据库")

	//_db, err := sql.Open("sqlite3", "./dns-srv._db")
	_db, err := sql.Open("sqlite3", os.ExpandEnv(itDns.Dsn))
	if err != nil {

		log.Error("打开数据库错误：", err)

		return nil, err
	}

	_db.SetConnMaxLifetime(itDns.MaxLifetime)

	_db.SetMaxOpenConns(itDns.MaxOpenConnections)

	_db.SetMaxIdleConns(itDns.MaxIdleConnections)

	return _db, nil
}

func (itDns *ItDns) close(db *sql.DB) {
	log.Info("=关闭数据库=>")
	err := db.Close()
	if err != nil {
		log.Error("=数据库异常=>", err)
	}
}
